'''
    prim是一种最小权重生成树算法：什么是最小权重生成树这里不介绍了。
    他和Dijkstra算法的区别：Dijkstra算法是指最短路径的算法，其distance中存储的是节点到根结点的最短路径。
                        Prim是最小权重生成树算法，其distance存储的是到上一个节点的最短距离。其算法流程比较简单，以前学习过。不好描述就不描述了
                        算法描述见文档
'''
from Dijkstra import *
import sys

# 这里要注意那个堆里面能否添加相同的元素，顶点相同时要覆盖原来的值。我在原来的堆中相当于添加了一个任何位置删除元素的方法。
def prim(G,start):
    pq = PriorityQueue()
    for v in G:
        v.setDistance(sys.maxsize)
        v.setPred(None)
    start.setDistance(0)
    start.setPred(None)
    pq.buildHeap([(v.getDistance(),v) for v in G])
    while not pq.isEmpty():
        currentVert = pq.delMin()
        # print(currentVert.getId())
        for nextVert in currentVert.getConnections():
          newCost = currentVert.getWeight(nextVert)
          if nextVert in pq and newCost<nextVert.getDistance():
              nextVert.setPred(currentVert)
              nextVert.setDistance(newCost)
              pq.decreaseKey(nextVert,newCost)

if __name__ == '__main__':


    graph_dict = {
        "A": {"B": 5, "C": 1},
        "B": {"A": 5, "C": 2, "D": 1},
        "C": {"A": 1, "B": 2, "D": 4, "E": 8},
        "D": {"B": 1, "C": 4, "E": 3, "F": 6},
        "E": {"C": 8, "D": 3},
        "F": {"D": 6},
    }
    g = Graph()  # 创建图
    # 创建图
    for v in graph_dict:
        for w in graph_dict[v]:
            g.addEdge(v,w,graph_dict[v][w])

    # 打印图：
    print('图创建如下:')
    for v in g: # 这个返回的是键值：顶点
        for w in v.getConnections():
            print('\t(%s,%s,%s)' %(v.getId(),w.getId(),v.connectTo[w]))

    prim(g,g.getVertex('A'))

    #打印出最短路径,需要反序打印，当然可以用栈
    print('最小生成树路径：')
    for v in g: # 遍历g中的顶点。
        h = v
        res = []
        res.append([h.getId()])

        while h.getPred() != None:
            h = h.getPred()
            # print(h.getId())
            res.append(h.getId())

        i = len(res)
        while i > 1:
            print(res[i-1],'->',end='')
            i = i - 1

        print(res[0][0])
